package com.tuniu.agents.common.message;

import com.tuniu.agents.message.MessageSender;
import com.tuniu.agents.message.StreamMessage;
import com.tuniu.agents.util.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;

@Slf4j
public class MessageSenderManager implements MessageSender {



    /**
     * sessionManager management
     */
    private final Map<String, FluxSessionManager<String>> conversationSessionManagers = new ConcurrentHashMap<>();

    /**
     * Store main responseId
     */
    private final Map<String, String> conversationResponseIds = new ConcurrentHashMap<>();


    public Flux<String> initConversation(Supplier<Flux<String>> fluxSupplier, String conversationId, String responseId) {
        // Create a new instance of FluxSessionManager.
        FluxSessionManager<String> sessionManager = new FluxSessionManager<>();
        conversationSessionManagers.put(conversationId,sessionManager);
        conversationResponseIds.put(conversationId,responseId);
        return sessionManager.initializeAndStartSession(fluxSupplier,responseId);
    }

    public void removeConversation(String conversationId) {
        conversationSessionManagers.remove(conversationId);
        conversationResponseIds.remove(conversationId);
    }

    /**
     * Cancel and stop conversation
     * @param conversationId Conversation ID
     * @return true if the conversation was successfully stopped, false otherwise
     */
    public boolean stopConversation(String conversationId) {
        // Clean up session resources on the current node
        if (!conversationSessionManagers.containsKey(conversationId)) {
            log.warn("Conversation {} does not exist, no need to clean up", conversationId);
            return false;
        }
        FluxSessionManager<String> sessionManager = conversationSessionManagers.get(conversationId);
        if (sessionManager != null) {
            // First cancel the underlying stream
            sessionManager.completeMainResponse();
        }
        // Then clean up resources
        removeConversation(conversationId);
        log.info("Conversation {} has been successfully stopped", conversationId);
        return true;
    }

    @Override
    public void send(String conversationId, String message) {
        FluxSessionManager<String> sessionManager = conversationSessionManagers.get(conversationId);
        if(sessionManager == null) {
            return;
        }
        String responseId = conversationResponseIds.get(conversationId);
        if (JsonUtil.isValidJson(message)){
            sessionManager.outputFlux(responseId, message);
        }else {
            sessionManager.outputFlux(responseId, BaseAgentNewResp.sendAnalysisMsg(message,1, responseId));
        }
    }

    @Override
    public void send(String conversationId, String responseId, String message) {
        FluxSessionManager<String> sessionManager = conversationSessionManagers.get(conversationId);
        if(sessionManager == null) {
            return;
        }
        if (JsonUtil.isValidJson(message)){
            sessionManager.outputFlux(responseId, message);
        }else{
            sessionManager.outputFlux(responseId, BaseAgentNewResp.sendAnswerMsg(responseId,message));
        }
    }

    @Override
    public void send(String conversationId, StreamMessage message) {
        FluxSessionManager<String> sessionManager = conversationSessionManagers.get(conversationId);
        if(sessionManager == null) {
            return;
        }
        String responseId = conversationResponseIds.get(conversationId);
        sessionManager.outputFlux(responseId,message.toEventString());
    }

    @Override
    public void send(String conversationId, String responseId, StreamMessage message) {
        FluxSessionManager<String> sessionManager = conversationSessionManagers.get(conversationId);
        if(sessionManager == null) {
            return;
        }
        sessionManager.outputFlux(responseId,message.toEventString());
    }

    @Override
    public void send(String conversationId, Flux<String> message) {
        FluxSessionManager<String> sessionManager = conversationSessionManagers.get(conversationId);
        if(sessionManager == null) {
            return;
        }
        String responseId = conversationResponseIds.get(conversationId);
        sessionManager.outputFlux(responseId,message);
    }

    @Override
    public void send(String conversationId, String responseId, Flux<String> message) {
        FluxSessionManager<String> sessionManager = conversationSessionManagers.get(conversationId);
        if(sessionManager == null) {
            return;
        }
        sessionManager.outputFlux(responseId,message);
    }


    /**
     * Get global variable
     * @param type Type of variable to retrieve
     * @return JSON representation of the requested global variable
     */
    public String getGlobalVariable(int type){
        if (type==1){
           return JsonUtil.toJson(conversationSessionManagers);
        }
        return  JsonUtil.toJson(conversationResponseIds);
    }
}
